# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: gekim

import calendar
import datetime
import json
import toml
import jinja2
import subprocess
import tempfile
import sys
import os
from absl import app, flags

FLAGS = flags.FLAGS

machines = [
    "feed-05.ap-northeast-1.aws",
    "feed-01.ap-northeast-1.aws.huobi",
    "feed-02.ap-northeast-2.aws",
    "feed-02.eu-west-1.aws",
    "feed-02.us-east-1.aws",
    "feed-05.cn-hongkong.aliyun",
]

week_futures = ('Futures.Okex.v3', 'Futures.Okex.v5', 'Futures.Huobi.v1')
quarter_futures = ('Futures.Bybit.v2', 'Futures.Ftx.v1', 'Futures.Binance.v1-delivery')


def get_last_friday_of_seasons(trading_date):
  year = trading_date.year
  seasons = [[3, 31], [6, 30], [9, 30], [12, 31]]
  quarter_days = []
  for day in seasons:
    m, d = day
    start_date = datetime.datetime(year, m, d)
    last_friday = start_date
    oneday = datetime.timedelta(days=1)
    while last_friday.weekday() != calendar.FRIDAY:
      last_friday -= oneday
    quarter_days.append(last_friday)
  return quarter_days


def build_mea_to_config_file(m, config, file_path):
  mea_to_config_file = {}
  for mea in config['feed']['exchanges'].keys():
    mea_to_config_file[(m, mea)] = file_path

  return mea_to_config_file

def search_vanilla_config(file_name, trading_date):
  vanillas_configs_dir = "../data/coin2/feed/vanillas_configs/"
  if not os.path.exists(vanillas_configs_dir):
    return None
  date_folders = []
  for _, dirs, _ in os.walk(vanillas_configs_dir, topdown=False):
    for name in dirs:
        if name.isnumeric():
            date_folders.append(name)
  date_folders.sort(reverse=True)
  for date_folder in date_folders:
    if int(trading_date) < int(date_folder):
      continue
    file_path = os.path.join(vanillas_configs_dir, date_folder, file_name)
    if os.path.exists(file_path):
      return file_path
  return None


def load_feed_writer_config(machines, trading_date_str):
  json_config_tmpl = "../data/coin2/feed/ops/feed-writer/%s.json"
  config = {}
  mea_to_config_file = {}
  for m in machines:
    json_config_path = search_vanilla_config("%s.json" % m, trading_date_str)
    json_config_path = json_config_tmpl % m
    with open(json_config_path, 'r') as f:
      json_obj = json.load(f)
      config[m] = json_obj
      config_file = json_config_path[3:]
      mea_to_config_file.update(build_mea_to_config_file(m, json_obj, config_file))
  return config, mea_to_config_file


def load_symbol_groups_setting(machine, trading_date_str):
  filter_out_meas = []
  symbol_groups_path = search_vanilla_config('%s_symbol_groups.toml' % machine, trading_date_str)
  if symbol_groups_path is None:
    symbol_groups_path = '../data/coin2/feed/symbol_groups.toml'
  with open(symbol_groups_path, 'r') as f:
    symbol_groups = toml.load(f)
    for mea in filter_out_meas:
      symbol_groups.pop(mea, None)
    return symbol_groups


def load_subscribers_setting():
  subscribers_path = "../data/coin2/feed/subscribers.toml"
  with open(subscribers_path, 'r') as f:
    return toml.load(f)


def convert_mea_to_toml_setting(mea):
  return mea.lower().replace('.', '_')


def get_mea_from_config(machine, config):
  return config[machine]['feed']['exchanges'].keys()


def get_recipe_from_config(machine, mea, config, subscribers):
  channels = config[machine]['feed']['exchanges'][mea]['channels']
  assert len(channels) > 0, "Only support subscribe by channel"
  toml_mea = convert_mea_to_toml_setting(mea)
  recipes_setting = subscribers[toml_mea]['recipes']
  recipes = []
  channels_set = set(channels)
  for recipe, channel_settings in recipes_setting.items():
    if channel_settings and set(channel_settings).issubset(channels_set):
      recipes.append(recipe)

  return recipes


def get_groups_from_setting(machine, mea, config):
  groups = config[machine]['feed']['exchanges'][mea]["products"]["groups"]
  assert len(groups) > 0, "Only support subscribe by group"
  return groups


def get_fastfeed_universe_setting(machine, mea):
  setting_path = 'data/coin2/feed/ops/fastfeed-universe-converter/common.json'
  return setting_path


use_data2_nfs_machines = [
    "feed-05.ap-northeast-1.aws",
    "feed-05.cn-hongkong.aliyun",
]


ignore_recipes = {
    'Futures.Binance.v1-delivery': ['snapshot'],
    'Futures.Binance.v1': ['snapshot'],
    'Futures.Huobi.v1': ['snapshot'],
    'Futures.Huobi.v1-swap': ['snapshot'],
    'Futures.Okex.v3': ['snapshot'],
    'Futures.Okex.v3-swap': ['snapshot'],
    'Futures.Okex.v5': ['snapshot'],
    'Futures.Okex.v5-swap': ['snapshot'],
}


def is_mea_enabled(mea, arbitration):
  return True


def gen_lauch_script(trading_date, settings, log_dir):
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)
  print(settings)
  template = template_env.get_template(
      'coin/support/feed_tool/feed_cache/script/coin2_universe.tmpl')
  return template.render(trading_date=trading_date, settings=settings, log_dir=log_dir)


def run_sbatch(script_filepath):
  sbatch_args = FLAGS.sbatch_args
  sbatch_args = [token for token in sbatch_args.split(' ') if token]
  # cwd setting let slurm-id.out log in fastfeed deploy directory
  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath],
                             stdout=subprocess.PIPE,
                             cwd='/remote/iosg/jenkins-2/bot-coin/jenkins/workspace/coin2_universe/')
  sbatch_output = completed.stdout.decode('utf-8')
  if completed.returncode != 0:
    print(sbatch_output + ' - FAILED')
    sys.exit(1)
  else:
    print(sbatch_output + ' - SUCCESS')


def main(argv):
  trading_date_str = FLAGS.trading_date
  proto_feed_universe_root = FLAGS.proto_feed_universe_root
  dest_proto_tag = FLAGS.dest_proto_tag
  assert trading_date_str, '--trading_date must be specified.'
  trading_date = datetime.datetime.strptime(trading_date_str, '%Y%m%d')

  config, mea_to_config_file = load_feed_writer_config(machines, trading_date_str)
  subscribers = load_subscribers_setting()
  quarter_days = get_last_friday_of_seasons(trading_date)
  job_id = 0
  settings = {}
  for m in machines:
    if (FLAGS.machine is not None) and (m not in FLAGS.machine):
      continue
    settings[m] = {}
    meas = get_mea_from_config(m, config)
    symbol_groups = load_symbol_groups_setting(m, trading_date_str)
    for mea in meas:
      if (FLAGS.mea is not None) and (mea not in FLAGS.mea):
        continue
      toml_mea = convert_mea_to_toml_setting(mea)
      norms = []
      for k,v in symbol_groups[toml_mea].items():
        norms += v
      norms = ",".join(norms)
      settings[m][mea] = {}
      recipes = get_recipe_from_config(m, mea, config, subscribers)
      filtered_recipes1 = []
      filtered_recipes2 = []  # Huobi futures non_tbs recipe need to be splitted.
      for recipe in recipes:
        if (FLAGS.recipe is not None) and (recipe not in FLAGS.recipe):
          continue
        if (FLAGS.recipe_include is not None) and (FLAGS.recipe_include not in recipe):
          continue
        if recipe in ignore_recipes.get(mea, []):
          continue
        if recipe == 'non_tbs' and mea == 'Futures.Huobi.v1':
          filtered_recipes2.append(recipe)
        else:
          filtered_recipes1.append(recipe)

      for filtered_recipes in (filtered_recipes1, filtered_recipes2):
        if len(filtered_recipes) == 0:
          continue
        recipes_str = ','.join(filtered_recipes)
        threads_num = len(filtered_recipes)
        # slurm22 will bind 2x cpus_per_task number of vcore to job due to hyperthreading
        # if threads_num <= 2:
        #   cpus_per_task = 1
        # else:
        #   cpus_per_task = 2

        settings[m][mea][recipes_str] = {}
        fastfeed_setting = get_fastfeed_universe_setting(m, mea)
        driver_settings = [fastfeed_setting, mea_to_config_file[(m, mea)]]
        start = trading_date_str
        settings[m][mea][recipes_str] = [
        (driver_settings, start, norms, job_id, proto_feed_universe_root, dest_proto_tag) ]
        #(driver_settings, start, threads_num, cpus_per_task, job_id, proto_feed_universe_root, dest_proto_tag)
        job_id = job_id + 1

  #log_dir = "/remote/iosg/jenkins-2/bot-coin/jenkins/workspace/coin2_fastfeed/logs"
  log_dir = FLAGS.log_dir
  script = gen_lauch_script(trading_date_str, settings, log_dir)

  if FLAGS.script_only:
    print(script)
    return

  tmp_dir = '/remote/iosg/jenkins-2/bot-coin/jenkins/workspace/coin2_universe/tmp'
  with tempfile.NamedTemporaryFile(mode='w', dir=tmp_dir) as script_file:
    script_file.write(script)
    script_file.flush()
    run_sbatch(script_file.name)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'yyyymmdd')

  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')

  flags.DEFINE_list('machine', None, "Separate by comma")

  flags.DEFINE_list('mea', None, "Separate by comma")

  flags.DEFINE_list('recipe', None, "Separate by comma")

  flags.DEFINE_string('recipe_include', None, "recipe substring")

  flags.DEFINE_list('group', None, "Separate by comma")

  flags.DEFINE_string('proto_feed_universe_root', None, '')

  flags.DEFINE_string('dest_proto_tag', None, '')

  flags.DEFINE_string(
      'sbatch_args',
      '-n150 --mem-per-cpu=1G --job-name=coin2_fastfeed --hint=multithread --priority=TOP --wait',
      'sbatch arguments')

  app.run(main)
